import os
import datetime
import pandas as pd
from openpyxl import Workbook, load_workbook
from openpyxl.styles import Alignment, Font, PatternFill, Border, Side
from openpyxl.utils import get_column_letter


def split_excel_by_header(
    input_dir="input",
    output_dir="output",
    header_name=None,
    extra_filters=None,
    file_name=None,
    split_to="files",  # 'files' 或 'sheets'
    save_func=None,
):
    """
    将 input_dir 中的一个 Excel 文件根据 header_name 列的值拆分

    参数:
    - split_to: 'files' 拆分到不同文件, 'sheets' 拆分到同一文件的不同工作表
    """
    os.makedirs(output_dir, exist_ok=True)

    # 找文件
    files = [f for f in os.listdir(input_dir) if f.lower().endswith((".xlsx", ".xls"))]
    if not files:
        raise ValueError("未找到 Excel 文件，请重新上传！")

    target_file = file_name or files[0]
    in_path = os.path.join(input_dir, target_file)
    if not os.path.exists(in_path):
        raise FileNotFoundError(f"未找到指定输入文件: {in_path}")

    # 读取 Excel
    try:
        df = pd.read_excel(in_path, dtype=str)
    except Exception as e:
        raise ValueError(f"读取 Excel 失败：{e}")

    if not header_name:
        raise ValueError("必须指定要分组的表头名称 (header_name)")

    if header_name not in df.columns:
        raise ValueError(f"指定的表头 '{header_name}' 在文件中不存在")

    # 标准化 extra_filters
    ef = {}
    if extra_filters:
        if isinstance(extra_filters, dict):
            for k, v in extra_filters.items():
                if v is None:
                    ef[k] = []
                elif isinstance(v, (list, tuple)):
                    ef[k] = [str(x).strip() for x in v]
                else:
                    ef[k] = [s.strip() for s in str(v).split(",")]
        else:
            raise ValueError('extra_filters 必须为 dict，形如 {"备注":["", "值1"]}')

    out_paths = []
    now = datetime.datetime.now()
    ym = f"{now.year}年{now.month}月"

    grp_vals = df[header_name].fillna("").astype(str).unique().tolist()

    if split_to == "sheets":
        # 拆分到同一文件的不同工作表
        output_file = f"细分数据_{ym}.xlsx"
        output_path = os.path.join(output_dir, output_file)

        try:
            wb = Workbook()
            # 删除默认的工作表
            wb.remove(wb.active)

            for gv in grp_vals:
                # 分组数据
                if gv == "" or gv is None:
                    mask = df[header_name].isna() | df[header_name].astype(str).map(
                        lambda s: str(s).strip() == ""
                    )
                else:
                    mask = (
                        df[header_name]
                        .astype(str)
                        .map(lambda s: str(s).strip() == str(gv).strip())
                    )

                sub = df[mask].copy()
                if sub.empty:
                    continue

                # 应用额外筛选
                for ex_col, ex_values in ef.items():
                    if ex_col not in sub.columns:
                        sub = pd.DataFrame()
                        break
                    if not ex_values:
                        continue

                    def match_cell(cell):
                        sval = "" if pd.isna(cell) else str(cell).strip()
                        for ev in ex_values:
                            if ev == "":
                                if sval == "":
                                    return True
                            elif sval == str(ev).strip():
                                return True
                        return False

                    sub = sub[sub[ex_col].apply(match_cell)]
                    if sub.empty:
                        break

                if sub.empty:
                    continue

                # 处理序号列
                if "序号" in sub.columns:
                    sub["序号"] = range(1, len(sub) + 1)
                    cols = ["序号"] + [c for c in sub.columns if c != "序号"]
                    sub = sub[cols]
                else:
                    sub.insert(0, "序号", range(1, len(sub) + 1))

                # 创建工作表
                safe_group = (
                    "".join(ch for ch in str(gv) if ch not in r'\/:*?"<>|').strip()
                    or "空值"
                )
                # 工作表名称长度限制
                sheet_name = safe_group[:31]
                ws = wb.create_sheet(sheet_name)

                # 写入数据并设置样式
                _write_data_to_sheet(ws, sub, safe_group, ym)

            if save_func:
                save_func(wb, output_path)
            else:
                wb.save(output_path)

            out_paths.append(output_path)

        except Exception as e:
            raise ValueError(f"保存文件失败：{e}")

    else:
        # 拆分到不同文件（原逻辑）
        for gv in grp_vals:
            # 分组数据
            if gv == "" or gv is None:
                mask = df[header_name].isna() | df[header_name].astype(str).map(
                    lambda s: str(s).strip() == ""
                )
            else:
                mask = (
                    df[header_name]
                    .astype(str)
                    .map(lambda s: str(s).strip() == str(gv).strip())
                )

            sub = df[mask].copy()
            if sub.empty:
                continue

            # 应用额外筛选
            for ex_col, ex_values in ef.items():
                if ex_col not in sub.columns:
                    sub = pd.DataFrame()
                    break
                if not ex_values:
                    continue

                def match_cell(cell):
                    sval = "" if pd.isna(cell) else str(cell).strip()
                    for ev in ex_values:
                        if ev == "":
                            if sval == "":
                                return True
                        elif sval == str(ev).strip():
                            return True
                    return False

                sub = sub[sub[ex_col].apply(match_cell)]
                if sub.empty:
                    break

            if sub.empty:
                continue

            # 处理序号列
            if "序号" in sub.columns:
                sub["序号"] = range(1, len(sub) + 1)
                cols = ["序号"] + [c for c in sub.columns if c != "序号"]
                sub = sub[cols]
            else:
                sub.insert(0, "序号", range(1, len(sub) + 1))

            # 输出文件名
            safe_group = (
                "".join(ch for ch in str(gv) if ch not in r'\/:*?"<>|').strip()
                or "空值"
            )
            out_name = f"{safe_group}-{ym}.xlsx"
            out_path = os.path.join(output_dir, out_name)

            try:
                wb = Workbook()
                ws = wb.active

                # 写入数据并设置样式
                _write_data_to_sheet(ws, sub, safe_group, ym)

                if save_func:
                    save_func(wb, out_path)
                else:
                    wb.save(out_path)
            except Exception as e:
                raise ValueError(f"保存文件失败：{e}")

            out_paths.append(out_path)

    if not out_paths:
        raise ValueError("没有匹配到任何符合条件的分组，未生成文件。")

    return out_paths


def _write_data_to_sheet(ws, data, group_name, year_month):
    """将数据写入工作表并设置样式"""
    # 标题行
    title = f"{group_name} - {year_month}"
    final_cols = list(data.columns)
    last_col_idx = len(final_cols)
    ws.cell(row=1, column=1, value=title)
    ws.merge_cells(start_row=1, start_column=1, end_row=1, end_column=last_col_idx)
    title_cell = ws.cell(row=1, column=1)
    title_cell.font = Font(bold=True, size=16, name="微软雅黑")
    title_cell.alignment = Alignment(
        horizontal="center", vertical="center", wrap_text=True
    )
    title_cell.fill = PatternFill(fill_type="solid", fgColor="D9D9D9")

    # 表头
    for ci, col_name in enumerate(final_cols, start=1):
        cell = ws.cell(row=2, column=ci, value=col_name)
        cell.font = Font(bold=True)
        cell.alignment = Alignment(
            horizontal="center", vertical="center", wrap_text=True
        )
        cell.fill = PatternFill(fill_type="solid", fgColor="F2F6FA")

    # 数据
    for r_idx, row in enumerate(data.itertuples(index=False, name=None), start=3):
        for c_idx, value in enumerate(row, start=1):
            cell = ws.cell(row=r_idx, column=c_idx, value=value)
            cell.alignment = Alignment(
                horizontal="center", vertical="center", wrap_text=True
            )

    ws.row_dimensions[1].height = 40
    for r in range(2, ws.max_row + 1):
        ws.row_dimensions[r].height = 25

    thin = Side(border_style="thin", color="000000")
    cell_border = Border(left=thin, right=thin, top=thin, bottom=thin)
    for c in range(1, last_col_idx + 1):
        for r in range(1, ws.max_row + 1):
            ws.cell(row=r, column=c).border = cell_border

    # 自适应列宽
    for idx, col in enumerate(final_cols, start=1):
        letter = get_column_letter(idx)
        try:
            max_len = max(
                data[col].fillna("").astype(str).map(len).max(), len(str(col))
            )
        except Exception:
            max_len = len(str(col))
        ws.column_dimensions[letter].width = min(max(max_len + 2, 8), 40)
